﻿import logging
import os
import re
import socket
import sys
import hashlib  # 哈希算法
import json
from ast import literal_eval
from getpass import getpass
from os.path import isfile, normpath,  join  # 文件路径处理
from struct import pack, unpack, calcsize  # 用于处理二进制数据打包/解包
from tqdm import tqdm  # 进度条显示

import rsa  # 添加RSA库
from cryptography.fernet import Fernet  # 添加对称加密
from tqdm import tqdm  # 进度条

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('ftp_client.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

intSize = calcsize('i')# 定义计算整数大小的常量（用于网络通信中长度前缀）
bufferSize = 40960


class FTPClient:
    def __init__(self, host='localhost', port=3306):
        """初始化FTP客户端"""
        self.host = host  # 服务器IP
        self.port = port  # 服务器端口
        self.socket = None  # Socket连接对象
        self.cipher = None  # 加密器（Fernet对称加密）
        self.current_dir = "."  # 当前目录（客户端本地）
        self.dir_stack = []  # 目录栈（用于pushd/popd功能）
        self.timeout = 10  # 默认超时时间（秒），减少等待时间
        self.socket_buffer_size = 8192  # Socket缓冲区大小
        self.max_retries = 3  # 最大重试次数
        self.username = None  # 保存用户名，用于重连
        self.password = None  # 保存密码，用于重连
        self.session_key = None  # 保存会话密钥
        self.private_key = None  # 保存私钥

    def connect(self, reconnect=False):
        """连接到服务器并完成认证流程"""
        try:
            # 如果是重连且已存在连接，先关闭
            if reconnect and self.socket:
                try:
                    self.socket.close()
                except:
                    pass
                self.socket = None
                
            # 1. 建立TCP连接
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.socket.settimeout(self.timeout)  # 设置超时时间
            self.socket.connect((self.host, self.port))
            print(f"已连接到服务器 {self.host}:{self.port}")
            
            # 如果是重连且已有登录信息，使用已保存的凭据
            if reconnect and self.username and self.session_key and self.cipher:
                print(f"正在使用已保存的凭据重新登录...")
                # 发送用户名
                self._send_data(self.username.encode())
                
                # 接收服务器响应（假设服务器会发送会话密钥）
                length = unpack('i', self.socket.recv(intSize))[0]
                encrypted_key = self.socket.recv(length)
                
                try:
                    # 使用已保存的私钥解密会话密钥
                    session_key = rsa.decrypt(encrypted_key, self.private_key)
                    self.cipher = Fernet(session_key)
                    logger.info("会话密钥解密成功")
                    
                    # 发送加密的凭据
                    cred = f"{self.username},{self.password}"
                    encrypted_cred = self.cipher.encrypt(cred.encode())
                    self._send_data(encrypted_cred)
                    
                    # 接收认证结果
                    length = unpack('i', self.socket.recv(intSize))[0]
                    encrypted_response = self.socket.recv(length)
                    response = self.cipher.decrypt(encrypted_response)
                    
                    if response == b'ok':
                        print("重新连接成功")
                        logger.info("重新连接成功")
                        return True
                    else:
                        print("重新连接失败，需要重新登录")
                        # 继续普通登录流程
                except Exception as e:
                    print(f"重新连接失败: {e}")
                    # 继续普通登录流程

            # 2. 获取用户名
            self.username = input("请输入用户名: ")
            
            # 3. 加载或生成RSA密钥对
            pub_key, priv_key = load_or_generate_keys(self.username)
            self.private_key = priv_key  # 保存私钥用于重连

            # 4. 发送用户名（服务器用此查找用户公钥）
            self._send_data(self.username.encode())

            # 5. 接收服务器响应
            length = unpack('i', self.socket.recv(intSize))[0]
            data = self.socket.recv(length)
            
            # 检查是否是公钥请求
            if length == 11 and data == b'need_pubkey':
                logger.info("服务器请求公钥，正在发送...")
                print("首次登录，正在设置公钥...")
                
                # 读取本地生成的公钥
                with open(f"{self.username}_public.pem", 'rb') as f:
                    public_key_data = f.read()
                
                # 发送公钥给服务器 (使用_send_data方法确保有正确的长度前缀)
                self._send_data(public_key_data)
                logger.info("已发送公钥到服务器")
                
                # 接收服务器加密的会话密钥
                length = unpack('i', self.socket.recv(intSize))[0]
                encrypted_key = self.socket.recv(length)
                logger.info("接收到加密会话密钥")
            else:
                encrypted_key = data
                logger.info("接收到加密会话密钥")

            # 6. 用RSA私钥解密出AES会话密钥
            try:
                session_key = rsa.decrypt(encrypted_key, priv_key)
                self.session_key = session_key  # 保存会话密钥用于重连
                self.cipher = Fernet(session_key)  # 初始化对称加密器
                logger.info("会话密钥解密成功")
            except rsa.DecryptionError as e:
                logger.error(f"解密会话密钥失败: {str(e)}")
                print(f"解密会话密钥失败。可能是密钥不匹配。")
                print("解决方案:")
                print("1. 删除本地密钥文件并重新登录，生成新密钥")
                print("2. 使用update_key命令更新服务器上的公钥")
                self.close()
                return False
            except Exception as e:
                logger.error(f"解密会话密钥失败: {str(e)}")
                print(f"解密会话密钥失败: {e}")
                self.close()
                return False

            # 7. 发送加密的用户凭证（用户名+密码）
            self.password = getpass('请输入密码: ')
            cred = f"{self.username},{self.password}"
            encrypted_cred = self.cipher.encrypt(cred.encode())
            self._send_data(encrypted_cred)

            # 8. 接收认证结果
            length = unpack('i', self.socket.recv(intSize))[0]
            encrypted_response = self.socket.recv(length)
            response = self.cipher.decrypt(encrypted_response)

            if response == b'error':
                print('用户名或密码错误')
                logger.warning("登录失败: 用户名或密码错误")
                self.close()
                return False
            elif response == b'key_mismatch':
                # 密钥不匹配，提示用户更新公钥
                print('登录密钥不匹配，需要更新公钥')
                logger.warning("登录失败: 密钥不匹配")
                
                update_choice = input('是否更新公钥? (y/n): ').strip().lower()
                if update_choice == 'y':
                    # 发送更新公钥的命令
                    encrypted_command = self.cipher.encrypt('update_key'.encode())
                    self._send_data(encrypted_command)
                    
                    try:
                        # 接收服务器响应
                        length = unpack('i', self.socket.recv(intSize))[0]
                        encrypted_response = self.socket.recv(length)
                        response = self.cipher.decrypt(encrypted_response).decode()
                        
                        if response == "ready_for_key":
                            # 读取本地公钥文件并发送
                            with open(f"{self.username}_public.pem", 'rb') as f:
                                public_key_data = f.read()
                            
                            encrypted_key = self.cipher.encrypt(public_key_data)
                            self._send_data(encrypted_key)
                            
                            # 接收更新结果
                            length = unpack('i', self.socket.recv(intSize))[0]
                            encrypted_result = self.socket.recv(length)
                            result = self.cipher.decrypt(encrypted_result).decode()
                            
                            if result == "success":
                                print("公钥已成功更新，请重新登录")
                            else:
                                print("公钥更新失败")
                    except Exception as e:
                        logger.error(f"更新公钥失败: {str(e)}")
                        print(f"更新公钥失败: {str(e)}")
                
                self.close()
                return False
            elif response != b'ok':
                print(f"认证失败: {self._get_error_message(response)}")
                self.close()
                return False
            else:
                print("认证成功")
                logger.info("登录成功")
                return True

        except Exception as e:
            print(f"连接错误: {e}")
            logger.error(f"连接错误: {e}")
            self.close()
            return False
            
    def reconnect(self):
        """尝试重新建立连接"""
        print("检测到连接问题，正在尝试重新连接...")
        if self.connect(reconnect=True):
            print("重新连接成功")
            return True
        else:
            print("重新连接失败")
            return False

    def _send_data(self, data):
        """发送带长度前缀的数据（防止粘包）"""
        try:
            # 格式：4字节长度 + 实际数据
            self.socket.sendall(pack('i', len(data)) + data)
        except socket.timeout:
            raise Exception("发送数据超时")
        except Exception as e:
            raise Exception(f"发送数据失败: {e}")

    def _send_command(self, command):
        """发送加密后的命令到服务器"""
        encrypted_cmd = self.cipher.encrypt(command.encode())
        self._send_data(encrypted_cmd)

    def _receive_response(self):
        """接收服务器响应并解密，支持重试"""
        retries = 0
        last_error = None
        
        while retries < self.max_retries:
            try:
                # 先读长度
                length_data = self.socket.recv(intSize)
                if not length_data or len(length_data) != intSize:
                    raise Exception("接收数据长度信息失败")
                    
                length = unpack('i', length_data)[0]
                
                # 检查长度是否合理
                if length <= 0 or length > 10 * 1024 * 1024:  # 限制最大响应为10MB
                    raise Exception(f"接收到无效的数据长度: {length}")
                
                # 再读数据
                data = bytearray()
                bytes_received = 0
                
                while bytes_received < length:
                    chunk_size = min(self.socket_buffer_size, length - bytes_received)
                    chunk = self.socket.recv(chunk_size)
                    
                    if not chunk:
                        raise Exception("连接中断")
                        
                    data.extend(chunk)
                    bytes_received += len(chunk)
                    
                # 解密返回
                return self.cipher.decrypt(bytes(data))
                
            except socket.timeout:
                retries += 1
                last_error = "接收响应超时"
                print(f"接收超时，正在重试... ({retries}/{self.max_retries})")
                if retries >= self.max_retries:
                    break
            except Exception as e:
                retries += 1
                last_error = str(e)
                print(f"接收错误: {e}，正在重试... ({retries}/{self.max_retries})")
                if retries >= self.max_retries:
                    break
                    
        # 重试失败，抛出异常
        raise Exception(f"接收响应失败，已重试{self.max_retries}次: {last_error}")

    def _parse_response(self, response):
        """解析服务器响应，处理字节和字符串格式"""
        if isinstance(response, bytes):
            try:
                return response.decode().strip()
            except:
                return response
        return response
        
    def _is_ok_response(self, response):
        """检查响应是否表示成功"""
        if isinstance(response, bytes):
            return response == b'ok' or self._parse_response(response).lower() == 'ok'
        return response.lower() == 'ok'
        
    def _get_error_message(self, response):
        """从响应中提取错误信息"""
        if isinstance(response, bytes):
            try:
                return response.decode()
            except:
                return "未知错误"
        return response

    def find_files(self, pattern):
        """查找匹配指定模式的文件"""
        self._send_command(f"find {pattern}")
        try:
            length = unpack('i', self.socket.recv(intSize))[0]
            encrypted_files = self.socket.recv(length)
            files = literal_eval(self.cipher.decrypt(encrypted_files).decode())
            
            if not files:
                print(f"没有找到匹配 '{pattern}' 的文件")
            else:
                print(f"找到 {len(files)} 个匹配文件:")
                print(*files, sep='\n')
        except Exception as e:
            print(f"查找文件失败: {e}")
            
    def search_files(self, text):
        """搜索文件内容"""
        self._send_command(f"search {text}")
        try:
            length = unpack('i', self.socket.recv(intSize))[0]
            encrypted_results = self.socket.recv(length)
            results = literal_eval(self.cipher.decrypt(encrypted_results).decode())
            
            if not results:
                print(f"没有找到包含 '{text}' 的文件")
            else:
                print(f"找到 {len(results)} 个包含搜索文本的文件:")
                for filename, matches in results.items():
                    print(f"\n文件: {filename}")
                    for line_num, line in matches:
                        print(f"  第 {line_num} 行: {line.strip()}")
        except Exception as e:
            print(f"搜索文件内容失败: {e}")

    def download_file(self, filename):
        """下载文件（支持覆盖和断点续传）"""
        # 检查本地是否已存在同名文件
        local_size = 0
        resume = False
        
        if isfile(filename):
            local_size = os.path.getsize(filename)
            if local_size > 0:
                choice = input(f"本地文件 {filename} 已存在({local_size}字节)，是否断点续传(r)、覆盖(y)还是取消(n)?: ").lower()
                if choice == 'n':
                    print("下载取消")
                    return
                elif choice == 'r':
                    resume = True
                    print(f"断点续传: 从 {local_size} 字节开始")
                # 默认为覆盖
            else:
                overwrite = input(f"本地文件 {filename} 已存在，是否覆盖？(y/n): ").lower()
                if overwrite != 'y':
                    print("下载取消")
                    return

        # 发送下载命令，附带断点续传信息
        if resume:
            self._send_command(f"get {filename} {local_size}")
        else:
            self._send_command(f"get {filename}")
        
        try:
            # 接收服务器响应
            response = self._receive_response()
            
            # 检查响应
            parsed_response = self._parse_response(response)
            if parsed_response.lower() == 'no':
                print("文件不存在或无法下载")
                return
            elif parsed_response.lower().startswith('file_size:'):
                # 服务器返回文件大小信息
                try:
                    total_size = int(parsed_response.split(':')[1])
                    
                    if total_size == local_size:
                        print(f"文件已完全下载: {filename} ({total_size} 字节)")
                        return
                    
                    if resume and total_size < local_size:
                        print("错误: 本地文件大于服务器文件，无法断点续传")
                        choice = input("是否覆盖下载? (y/n): ").lower()
                        if choice != 'y':
                            print("下载取消")
                            return
                        resume = False
                        # 重新发送下载命令，不带断点续传信息
                        self._send_command(f"get {filename}")
                        response = self._receive_response()
                        
                    # 发送确认继续下载
                    self._send_command("continue")
                    
                    # 接收文件数据
                    file_mode = 'ab' if resume else 'wb'
                    received_size = local_size if resume else 0
                    
                    with open(filename, file_mode) as f:
                        # 创建进度条
                        progress = tqdm(
                            initial=received_size,
                            total=total_size,
                            unit='B',
                            unit_scale=True,
                            desc=f"下载 {filename}",
                            ncols=75
                        )
                        
                        # 开始接收数据
                        while received_size < total_size:
                            # 接收块大小信息
                            try:
                                chunk_size_data = self.socket.recv(intSize)
                                if not chunk_size_data:
                                    print("连接中断")
                                    break
                                    
                                chunk_size = unpack('I', chunk_size_data)[0]
                                
                                # 接收加密数据块
                                encrypted_chunk = bytearray()
                                bytes_received = 0
                                
                                while bytes_received < chunk_size:
                                    buffer_size = min(self.socket_buffer_size, chunk_size - bytes_received)
                                    chunk = self.socket.recv(buffer_size)
                                    
                                    if not chunk:
                                        raise Exception("接收数据块时连接中断")
                                        
                                    encrypted_chunk.extend(chunk)
                                    bytes_received += len(chunk)
                                
                                # 解密数据块并写入文件
                                decrypted_chunk = self.cipher.decrypt(bytes(encrypted_chunk))
                                f.write(decrypted_chunk)
                                
                                # 更新进度
                                chunk_actual_size = len(decrypted_chunk)
                                received_size += chunk_actual_size
                                progress.update(chunk_actual_size)
                                
                            except socket.timeout:
                                print("\n接收数据超时，下载中断")
                                break
                            except Exception as e:
                                print(f"\n接收数据时出错: {e}")
                                break
                        
                        progress.close()
                        
                    # 检查是否完成下载
                    if received_size >= total_size:
                        print(f"文件下载完成: {filename} ({total_size} 字节)")
                    else:
                        print(f"下载未完成: 已接收 {received_size}/{total_size} 字节")
                        print("下次可尝试断点续传")
                        
                except ValueError:
                    print(f"服务器返回了无效的文件大小信息: {parsed_response}")
                except Exception as e:
                    print(f"下载过程中出错: {e}")
            elif not self._is_ok_response(response):
                print(f"下载失败: {self._get_error_message(response)}")
                
        except Exception as e:
            print(f"下载命令执行失败: {e}")

    def list_files(self):
        """获取并显示当前目录文件列表"""
        self._send_command("list")
        try:
            response = self._receive_response()
            decoded_data = self._parse_response(response)
            
            # 尝试不同的方式解析服务器返回的数据
            try:
                # 首先尝试使用literal_eval解析
                files = literal_eval(decoded_data)
                if not files:
                    print("目录为空")
                else:
                    print("目录内容:")
                    for item in files:
                        print(item)
            except (SyntaxError, ValueError):
                # 如果解析失败，直接显示原始数据
                print("目录内容:")
                print(decoded_data)
        except Exception as e:
            print(f"获取文件列表失败: {e}")
            
    def change_dir(self, directory):
        """切换到指定目录"""
        self._send_command(f"cd {directory}")
        try:
            response = self._receive_response()
            
            if self._is_ok_response(response):
                print(f"已切换到目录 {directory}")
            else:
                print(f"切换目录失败: {self._get_error_message(response)}")
        except Exception as e:
            print(f"切换目录失败: {e}")
            
    def parent_dir(self):
        """切换到上级目录"""
        self._send_command("cd..")
        try:
            response = self._receive_response()
            
            if self._is_ok_response(response):
                print("已切换到上级目录")
            else:
                print(f"切换到上级目录失败: {self._get_error_message(response)}")
        except Exception as e:
            print(f"切换到上级目录失败: {e}")
            
    def show_current_dir(self):
        """显示当前工作目录"""
        self._send_command("pwd")
        try:
            response = self._receive_response()
            print(self._parse_response(response))
        except Exception as e:
            print(f"获取当前目录失败: {e}")

    def upload_file(self, filename):
        """上传文件到服务器，简化版本，减少通信次数"""
        if not isfile(filename):
            print(f"文件 {filename} 不存在")
            return
            
        file_size = os.path.getsize(filename)
        if file_size == 0:
            print("警告: 文件大小为0字节")
            
        print(f"准备上传: {filename} ({file_size} 字节)")
        
        try:
            # 简化上传过程：命令和文件名作为一个命令发送
            self._send_command(f"put {filename}")
            
            # 检查服务器是否接受上传
            try:
                response = self._receive_response()
                if not self._is_ok_response(response):
                    print(f"服务器拒绝上传: {self._get_error_message(response)}")
                    return
            except Exception as e:
                print(f"与服务器通信失败: {e}")
                return
                
            # 读取并发送文件
            try:
                # 直接发送文件大小（不加密）
                self.socket.sendall(pack('I', file_size))
                
                # 使用进度条显示上传进度
                with tqdm(total=file_size, unit='B', unit_scale=True, desc=f"上传 {filename}") as pbar:
                    bytes_sent = 0
                    
                    # 分块读取和发送文件
                    with open(filename, 'rb') as f:
                        while bytes_sent < file_size:
                            # 读取一块数据（较小块，避免超时）
                            chunk_size = min(16384, file_size - bytes_sent)  # 16KB chunks
                            chunk = f.read(chunk_size)
                            if not chunk:
                                break
                                
                            # 加密并发送
                            encrypted_chunk = self.cipher.encrypt(chunk)
                            self.socket.sendall(encrypted_chunk)
                            
                            # 更新进度
                            bytes_sent += len(chunk)
                            pbar.update(len(chunk))
                
                # 接收上传结果（设置较短的超时时间）
                original_timeout = self.socket.gettimeout()
                self.socket.settimeout(5)  # 设置短超时时间
                
                try:
                    result_response = self._receive_response()
                    if self._is_ok_response(result_response):
                        print(f"文件 {filename} 上传成功！")
                    else:
                        print(f"上传失败: {self._get_error_message(result_response)}")
                except Exception as e:
                    print(f"上传状态未知: {e}")
                    # 检测连接状态
                    try:
                        # 发送简单的命令测试连接
                        self._send_command("ping")
                        self._receive_response()
                    except:
                        print("连接已中断，正在重新连接...")
                        self.reconnect()
                        
                # 恢复原超时时间
                self.socket.settimeout(original_timeout)
                
            except socket.timeout:
                print("上传超时，请检查网络连接")
                self.reconnect()
            except Exception as e:
                print(f"上传过程中出错: {e}")
                self.reconnect()
        except Exception as e:
            print(f"上传命令执行失败: {e}")
            self.reconnect()
            
    def delete_file(self, filename):
        """删除服务器上的文件"""
        self._send_command(f"del {filename}")
        try:
            response = self._receive_response()
            
            if self._is_ok_response(response):
                print(f"文件 {filename} 删除成功")
            else:
                print(f"文件 {filename} 删除失败: {self._get_error_message(response)}")
        except Exception as e:
            print(f"删除文件失败: {e}")
            
    def make_dir(self, dirname):
        """创建目录"""
        self._send_command(f"mkdir {dirname}")
        try:
            response = self._receive_response()
            
            if self._is_ok_response(response):
                print(f"目录 {dirname} 创建成功")
            else:
                print(f"目录 {dirname} 创建失败: {self._get_error_message(response)}")
        except Exception as e:
            print(f"创建目录失败: {e}")
            
    def remove_dir(self, dirname):
        """删除目录"""
        self._send_command(f"rmdir {dirname}")
        try:
            response = self._receive_response()
            
            if self._is_ok_response(response):
                print(f"目录 {dirname} 删除成功")
            else:
                print(f"目录 {dirname} 删除失败: {self._get_error_message(response)}")
        except Exception as e:
            print(f"删除目录失败: {e}")
            
    def rename(self, old_name, new_name):
        """重命名文件"""
        self._send_command(f"rename {old_name} {new_name}")
        try:
            response = self._receive_response()
            
            if self._is_ok_response(response):
                print(f"重命名成功: {old_name} -> {new_name}")
            else:
                print(f"重命名失败: {self._get_error_message(response)}")
        except Exception as e:
            print(f"重命名失败: {e}")
            
    def file_attributes(self, filename):
        """查看文件属性"""
        self._send_command(f"attr {filename}")
        try:
            response = self._receive_response()
            decoded_data = self._parse_response(response)
            
            # 尝试解析文件属性
            try:
                attrs = literal_eval(decoded_data)
                
                print("文件属性:")
                print(f"大小: {attrs['size']} 字节")
                print(f"修改时间: {attrs['modified']}")
                print(f"权限: {attrs['mode']}")
            except (SyntaxError, ValueError, KeyError):
                # 如果解析失败，直接显示原始数据
                if decoded_data.lower().startswith('error'):
                    print(f"获取文件属性失败: {decoded_data}")
                else:
                    print("文件属性:")
                    print(decoded_data)
        except Exception as e:
            print(f"获取文件属性失败: {e}")
            
    def path_type(self, path):
        """判断路径类型"""
        self._send_command(f"di {path}")
        try:
            response = self._receive_response()
            print(f"路径类型: {self._parse_response(response)}")
        except Exception as e:
            print(f"判断路径类型失败: {e}")
            
    def create_file(self, filename):
        """创建空文件"""
        self._send_command(f"create {filename}")
        try:
            response = self._receive_response()
            
            if self._is_ok_response(response):
                print(f"文件 {filename} 创建成功")
            else:
                print(f"文件 {filename} 创建失败: {self._get_error_message(response)}")
        except Exception as e:
            print(f"创建文件失败: {e}")
            
    def show_tree(self):
        """显示目录树"""
        self._send_command("tree")
        try:
            length = unpack('i', self.socket.recv(intSize))[0]
            encrypted_tree = self.socket.recv(length)
            tree = self.cipher.decrypt(encrypted_tree).decode()
            
            print("\n目录树:")
            print(tree)
        except Exception as e:
            print(f"获取目录树失败: {e}")
            
    def push_dir(self):
        """保存当前目录"""
        self._send_command("pushd")
        try:
            length = unpack('i', self.socket.recv(intSize))[0]
            encrypted_dir = self.socket.recv(length)
            saved_dir = self.cipher.decrypt(encrypted_dir).decode()
            
            print(f"已保存当前目录: {saved_dir}")
        except Exception as e:
            print(f"保存目录失败: {e}")
            
    def pop_dir(self):
        """恢复保存的目录"""
        self._send_command("popd")
        try:
            length = unpack('i', self.socket.recv(intSize))[0]
            encrypted_response = self.socket.recv(length)
            response = self.cipher.decrypt(encrypted_response)
            
            if response == b'ok':
                print("恢复目录成功")
            else:
                print("恢复目录失败")
        except Exception as e:
            print(f"恢复目录失败: {e}")
            
    def copy_file(self, source, destination):
        """复制文件"""
        self._send_command(f"copy {source} {destination}")
        try:
            length = unpack('i', self.socket.recv(intSize))[0]
            encrypted_response = self.socket.recv(length)
            response = self.cipher.decrypt(encrypted_response)
            
            if response == b'ok':
                print(f"复制成功: {source} -> {destination}")
            else:
                print("复制失败")
        except Exception as e:
            print(f"复制文件失败: {e}")
            
    def copy_dir(self, source, destination):
        """递归复制目录"""
        self._send_command(f"xcopy {source} {destination}")
        try:
            length = unpack('i', self.socket.recv(intSize))[0]
            encrypted_response = self.socket.recv(length)
            response = self.cipher.decrypt(encrypted_response)
            
            if response == b'ok':
                print(f"递归复制成功: {source} -> {destination}")
            else:
                print("递归复制失败")
        except Exception as e:
            print(f"递归复制失败: {e}")
            
    def move_file(self, source, destination):
        """移动文件"""
        self._send_command(f"move {source} {destination}")
        try:
            length = unpack('i', self.socket.recv(intSize))[0]
            encrypted_response = self.socket.recv(length)
            response = self.cipher.decrypt(encrypted_response)
            
            if response == b'ok':
                print(f"移动成功: {source} -> {destination}")
            else:
                print("移动失败")
        except Exception as e:
            print(f"移动文件失败: {e}")
            
    def send_chat(self, message):
        """发送聊天消息"""
        self._send_command(f"chat {message}")
        try:
            length = unpack('i', self.socket.recv(intSize))[0]
            encrypted_response = self.socket.recv(length)
            chat_history = self.cipher.decrypt(encrypted_response).decode()
            
            print("\n=== 聊天历史 ===")
            print(chat_history)
            print("================\n")
        except Exception as e:
            print(f"发送聊天消息失败: {e}")
            
    def add_user(self, username, password, home_dir, pubkey_path=None):
        """添加新用户（需要管理员权限）"""
        # 构造add命令
        command = f"add {username} {password} {home_dir}"
        
        # 如果提供了公钥路径，添加到命令中
        if pubkey_path and os.path.exists(pubkey_path):
            with open(pubkey_path, 'rb') as f:
                pubkey_data = f.read()
            # 将公钥编码为base64字符串（避免二进制传输问题）
            import base64
            pubkey_b64 = base64.b64encode(pubkey_data).decode('utf-8')
            command = f"{command} {pubkey_b64}"
        elif pubkey_path:  # 如果指定了路径但文件不存在
            print(f"警告: 公钥文件 {pubkey_path} 不存在")
        
        # 发送添加用户命令
        self._send_command(command)
        
        try:
            # 接收服务器响应
            response = self._receive_response()
            
            if self._is_ok_response(response):
                print(f"用户 {username} 添加成功")
            elif self._parse_response(response).lower() == 'permission_denied':
                print("权限不足，需要管理员权限")
            else:
                print(f"添加用户失败: {self._get_error_message(response)}")
        except Exception as e:
            print(f"添加用户失败: {e}")
            
    def update_key(self):
        """更新服务器上的用户公钥"""
        self._send_command("update_key")
        try:
            # 接收服务器响应
            length = unpack('i', self.socket.recv(intSize))[0]
            encrypted_response = self.socket.recv(length)
            response = self.cipher.decrypt(encrypted_response).decode()
            
            if response == "ready_for_key":
                # 读取本地公钥文件并发送
                # 用户名从连接时输入，这里假设我们知道用户名
                username = input("请确认用户名: ")
                try:
                    with open(f"{username}_public.pem", 'rb') as f:
                        public_key_data = f.read()
                        
                    encrypted_key = self.cipher.encrypt(public_key_data)
                    self._send_data(encrypted_key)
                    
                    # 接收更新结果
                    length = unpack('i', self.socket.recv(intSize))[0]
                    encrypted_result = self.socket.recv(length)
                    result = self.cipher.decrypt(encrypted_result).decode()
                    
                    if result == "success":
                        print("公钥已成功更新到服务器")
                    else:
                        print("公钥更新失败")
                except FileNotFoundError:
                    print(f"找不到公钥文件 {username}_public.pem")
            else:
                print("服务器拒绝接收公钥更新")
        except Exception as e:
            print(f"更新公钥失败: {e}")

    def close(self):
        """关闭连接并清理资源"""
        if self.socket:
            self.socket.close()
            self.socket = None
        print("连接已关闭")

    def _is_connected(self):
        """检查是否与服务器连接正常"""
        if not self.socket:
            return False
        
        try:
            # 设置非常短的超时时间
            original_timeout = self.socket.gettimeout()
            self.socket.settimeout(2)
            
            try:
                # 发送简单命令
                self._send_command("ping")
                response = self._receive_response()
                
                # 恢复原超时时间
                self.socket.settimeout(original_timeout)
                return True
            except:
                # 恢复原超时时间
                self.socket.settimeout(original_timeout)
                return False
        except:
            return False

    def run(self):
        """主交互循环"""
        if not self.connect():  # 先建立连接
            return

        try:
            while True:
                try:
                    # 检查连接状态
                    if not self._is_connected():
                        print("检测到连接已断开，正在重新连接...")
                        if not self.reconnect():
                            print("无法重新连接到服务器，请检查网络或服务器状态")
                            break
                    
                    command = input(f"ftp> ").strip()
                    if not command:
                        continue

                    # 退出命令
                    if command.lower() in ('quit', 'exit', 'q'):
                        break

                    try:
                        # 处理各种命令（实际调用对应方法）
                        if command.lower() in ('list', 'ls', 'dir'):
                            self.list_files()
                        elif command.lower().startswith('get '):
                            self.download_file(command.split(' ', 1)[1])
                        elif command.lower().startswith('put '):
                            self.upload_file(command.split(' ', 1)[1])
                        elif command.lower().startswith('del '):
                            self.delete_file(command.split(' ', 1)[1])
                        elif command.lower().startswith('md ') or command.lower().startswith('mkdir '):
                            self.make_dir(command.split(' ', 1)[1])
                        elif command.lower().startswith('rd ') or command.lower().startswith('rmdir '):
                            self.remove_dir(command.split(' ', 1)[1])
                        elif command.lower().startswith('rename ') or command.lower().startswith('ren '):
                            parts = command.split(' ', 2)
                            if len(parts) == 3:
                                self.rename(parts[1], parts[2])
                            else:
                                print("用法: rename/ren <old> <new>")
                        elif command.lower().startswith('attr ') or command.lower().startswith('attrib '):
                            self.file_attributes(command.split(' ', 1)[1])
                        elif command.lower().startswith('di '):
                            self.path_type(command.split(' ', 1)[1])
                        elif command.lower().startswith('create '):
                            self.create_file(command.split(' ', 1)[1])
                        elif command.lower() == 'tree':
                            self.show_tree()
                        elif command.lower() == 'pushd':
                            self.push_dir()
                        elif command.lower() == 'popd':
                            self.pop_dir()
                        elif command.lower().startswith('copy '):
                            parts = command.split(' ', 2)
                            if len(parts) == 3:
                                self.copy_file(parts[1], parts[2])
                            else:
                                print("用法: copy <src> <dst>")
                        elif command.lower().startswith('xcopy '):
                            parts = command.split(' ', 2)
                            if len(parts) == 3:
                                self.copy_dir(parts[1], parts[2])
                            else:
                                print("用法: xcopy <src> <dst>")
                        elif command.lower().startswith('move '):
                            parts = command.split(' ', 2)
                            if len(parts) == 3:
                                self.move_file(parts[1], parts[2])
                            else:
                                print("用法: move <src> <dst>")
                        elif command.lower().startswith('chmod '):
                            parts = command.split(' ')
                            if len(parts) == 3:
                                self.chmod(parts[1], parts[2])
                            else:
                                print("用法: chmod <file> <mode>")
                        elif command.lower().startswith('chat '):
                            self.send_chat(command.split(' ', 1)[1])
                        elif command.lower().startswith('find '):
                            self.find_files(command.split(' ', 1)[1])
                        elif command.lower().startswith('search '):
                            self.search_files(command.split(' ', 1)[1])
                        elif command.lower() == 'reconnect':
                            # 添加重连命令
                            if self.reconnect():
                                print("已重新连接到服务器")
                            else:
                                print("重新连接失败")
                        elif command.lower() == 'ping':
                            # 添加ping命令检测连接状态
                            if self._is_connected():
                                print("连接正常")
                            else:
                                print("连接已断开")
                        elif command.lower().startswith('add '):
                            # 处理添加用户命令
                            parts = command.split(' ')
                            if len(parts) >= 4:
                                # 获取基本参数
                                username = parts[1]
                                password = parts[2]
                                home_dir = parts[3]
                                
                                # 可选的公钥路径
                                pubkey_path = parts[4] if len(parts) > 4 else None
                                
                                self.add_user(username, password, home_dir, pubkey_path)
                            else:
                                print("用法: add <user> <pass> <home> [pubkey]")
                        elif command.lower() == 'update_key':
                            self.update_key()
                        elif command.lower() in ('cd', 'pwd', 'cwd'):
                            self.show_current_dir()
                        elif ''.join(command.split()) == 'cd..':
                            self.parent_dir()
                        elif command.lower().startswith('cd '):
                            self.change_dir(command.split(' ', 1)[1])
                        elif command.lower() in ('help', '?'):
                            self._show_help()
                        else:
                            print(f"未知命令: {command}")
                    except socket.timeout:
                        print("命令执行超时，尝试重新连接...")
                        self.reconnect()
                    except (ConnectionError, ConnectionResetError, ConnectionAbortedError, ConnectionRefusedError) as e:
                        print(f"连接错误: {e}，尝试重新连接...")
                        self.reconnect()
                    except Exception as e:
                        print(f"命令执行错误: {e}")
                        if "接收响应失败" in str(e) or "连接中断" in str(e):
                            print("检测到连接问题，尝试重新连接...")
                            self.reconnect()
                except KeyboardInterrupt:
                    print("\n操作已取消，按Ctrl+C再次中断将退出程序")
                    try:
                        input("按Enter继续...")
                    except KeyboardInterrupt:
                        print("\n退出程序...")
                        break
        except KeyboardInterrupt:
            print("\n退出程序...")
        finally:
            self.close()
            
    def _show_help(self):
        """显示帮助信息"""
        help_text = """
可用命令:
list/ls/dir    - 列出当前目录内容
cd..           - 切换到上级目录
cd/cwd/pwd     - 显示当前目录
cd <dir>       - 切换到指定目录
get <file>     - 下载文件
put <file>     - 上传文件
del <file>     - 删除文件
md/mkdir <dir> - 创建目录
rd/rmdir <dir> - 删除目录
rename/ren <old> <new> - 重命名文件
attr/attrib <file> - 查看文件属性
chmod <file> <mode> - 修改文件权限
di <path>      - 判断路径类型
create <file>  - 创建空文件
tree           - 显示目录树
pushd          - 保存当前目录
popd           - 恢复保存的目录
copy <src> <dst> - 复制文件
xcopy <src> <dst> - 递归复制目录
move <src> <dst> - 移动文件
find <pattern> - 查找匹配的文件
search <text>  - 搜索文件内容
chat <message> - 发送聊天消息
add <user> <pass> <home> [pubkey] - 添加用户(需管理员权限)
update_key     - 更新服务器上的用户公钥
reconnect      - 重新连接服务器
ping           - 检测连接状态
quit/exit/q    - 退出
help/?         - 显示帮助
"""
        print(help_text)

    def chmod(self, filename, mode):
        """修改文件权限"""
        self._send_command(f"chmod {filename} {mode}")
        try:
            length = unpack('i', self.socket.recv(intSize))[0]
            encrypted_response = self.socket.recv(length)
            response = self.cipher.decrypt(encrypted_response)
            
            if response == b'ok':
                print(f"文件 {filename} 权限已修改为 {mode}")
            else:
                print(f"修改文件权限失败: {response.decode()}")
        except Exception as e:
            print(f"修改文件权限失败: {e}")


def load_or_generate_keys(username):
    """加载或生成用户的RSA密钥对"""
    private_key_file = f"{username}_private.pem"
    public_key_file = f"{username}_public.pem"

    # 如果密钥文件不存在，生成新密钥对
    if not os.path.exists(private_key_file) or not os.path.exists(public_key_file):
        print(f"为 {username} 生成新的RSA密钥对...")
        (pub_key, priv_key) = rsa.newkeys(512)

        # 保存私钥
        with open(private_key_file, 'wb') as f:
            f.write(priv_key.save_pkcs1())

        # 保存公钥
        with open(public_key_file, 'wb') as f:
            f.write(pub_key.save_pkcs1())

        print(f"密钥对已保存到 {private_key_file} 和 {public_key_file}")
        print("注意：新生成的密钥与服务器可能不匹配。请联系管理员更新您的公钥。")
        print("或者删除新生成的密钥文件，重新使用系统分配的密钥。")
        return pub_key, priv_key

    # 加载现有密钥
    with open(private_key_file, 'rb') as f:
        priv_key = rsa.PrivateKey.load_pkcs1(f.read())

    with open(public_key_file, 'rb') as f:
        pub_key = rsa.PublicKey.load_pkcs1(f.read())

    print(f"从 {private_key_file} 和 {public_key_file} 加载密钥对")
    return pub_key, priv_key


if __name__ == '__main__':
    # 需要在命令提示符环境执行程序
    if len(sys.argv) != 2:
        print(f'Usage: {sys.argv[0]} serverIPAddress')
        print('Example: python 3ftpClient.py 127.0.0.1')
        exit()
        
    serverIP = sys.argv[1]
    # 使用正则表达式检查IP地址格式
    if re.match(r'^\d{1,3}.\d{1,3}.\d{1,3}.\d{1,3}$', serverIP):
        # 创建FTP客户端并运行
        client = FTPClient(host=serverIP, port=10800)
        client.run()
    else:
        print('服务器地址不合法')
        exit()
